# -*- coding: utf-8 -*-
from __future__ import division
import nltk, datetime, string
from nltk.corpus.reader.toolbox import *
from config import *
from xml.etree.ElementTree import *

class MyToolboxCorpusReader (ToolboxCorpusReader):
    """Implements an extension of the standard toolbox corpus reader. In a way this class provides
    additional tools that can be used to access the data in a toolbox lexicon.
    Provides:
     - saving of toolbox data (lexicon as Element) as xml, html or again toolbox;
     - a tool to select a subset of the data.
    """
        
    def output(self, lexicon, format='text', target=''):
        """Reads a toolbox file and returns or saves the contents in a requested format. Results can be saved
        or printed on the screen. A regex filter and/or a lemmatizer can be applied to the lemmas.
        @param lexicon: lexicon of type Element
        @param format: the return format (xml, html or text)
        @param target: the target file (default to screen)
        """
        # open target (or use std.output)
        try:
            output = codecs.open(target, "w", ENCODING)
        except:
            output = sys.stdout
        
        # depending on output format: format entries and save/print
        if string.lower(format)=='text': 
            output.write(self.formatAsText(self.__addHeader(self.sort(lexicon))))
        elif string.lower(format)=='xml':
            self.formatAsXML(self.__addHeader(self.sort(lexicon))).write(output)
        elif string.lower(format)=='html':
            self.formatAsHTML(self.__addHeader(self.sort(lexicon))).write(output)
        else:
            print 'unknown output format requested...'
        
        # close target if necessary
        if not output==sys.stdout: output.close()
 
    def __addHeader(self, lexicon):
        """Adds a header to a lexicon.
        @param lexicon: the lexicon of type Element
        """
        header = lexicon.find('header')
        if header.find('date') is None:
            date = SubElement(header,'date')
            date.text = datetime.datetime.now().strftime("%d/%m/%Y %H:%M")
        else:
            header.find('date').text = datetime.datetime.now().strftime("%d/%m/%Y %H:%M")
        return lexicon
 
    #---- selection ----
            
    def filter(self, lexicon, regex='.*', lemmatize=0):
        """Filters a lexicon file and returns an Element that has been filtered. Lemmatization
        of the lemma is optional, the regex specifies the filter. 
        @param lexicon: lexicon of type Element
        @param regex: the filter (regular expression)
        @param lemmatize: flag, if 1 lemma is lemmatized
        @rtype: Element
        """
        filteredLexicon = XML('<toolbox_data><header /></toolbox_data>')
        for e in lexicon.getiterator('record'):
            lemma = e.find('lx').text
            if lemmatize==1: 
                lemma=nltk.PorterStemmer().stem(lemma)
            if re.compile(regex).match(lemma):
                filteredLexicon.append(e)
        return filteredLexicon
      
    #---- format ----
        
    def formatAsXML(self, lexicon):
        """Formats an Element as an ElementTree, that can be written to the target.
        @param lexicon: Element containing the data
        @rtype: ElementTree
        """
        return ElementTree(lexicon)

    def formatAsHTML(self, lexicon):
        """Converts an Element to HTML and returns it as an ElementTree, that can 
        be written to the target.
        @param lexicon: Element containing the data
        @rtype: ElementTree
        """
        html = XML('<html><header><meta http-equiv="Content-type" content="text/html;charset=UTF-8" /><title>Lexicon</title></header><body /></html>')
        body = html.find('body')
        
        meta = SubElement(body, 'div')
        for item in lexicon.find('header').getchildren():
            field = SubElement(meta, 'div')
            field.text = '%s: %s' % (item.tag, item.text)
        
        data = SubElement(body, 'div')
        for record in lexicon.getiterator('record'):
            group = SubElement(data, 'div')
            field = SubElement(field, 'h3')
            field.text = record.find('lx').text          
            for item in record.getchildren():
                field = SubElement(group, 'div')
                if not item.tag == 'lx': field.text = '%s: %s' % (item.tag, item.text)
        return ElementTree(html)
    
    def formatAsText(self, lexicon):
        """Converts an Element to text, that can be written to the target.
        @param lexicon: Element containing the data
        @rtype: string
        """
        parsedLexicon = [['\\'+j.tag+' '+j.text for j in i] for i in lexicon]
        return "\n\n".join(["\n".join(i) for i in parsedLexicon])    
    
    #---- sorting ----
    def sort(self, lexicon):
        """Sorts a lexicon based on the lx field. Leaves the header untouched.
        @param lexicon: the lexicon (Element)
        """
        data = []
        for elem in lexicon.getiterator('record'):
            key = elem.find('lx').text
            data.append((key, elem))       
        data.sort()
        
        lexicon[1:] = [item[-1] for item in data]
        return lexicon
    
class MyToolboxCorpusAnalysis(ToolboxCorpusReader):
    """Mixes in the statistics functions for the extended reader.
    Provides:
     - quantitative analysis tools.
    """
    
    def numEntries(self, lexicon):
        """Calculates the number of entries in a lexicon.
        @param lexicon: the lexicon in Element format
        """
        return len(lexicon.findall('record'))
    
    def ratioVerbsScf(self, lexicon):
        """Calculates the average number of verbs per scf in a lexicon.
        @param lexicon: the lexicon in Element format
        """
        return self.numEntries(lexicon) / self.numScfs(lexicon)
    
    def ratioScfsVerb(self, lexicon):
        """Calculates the average number of scfs per verb in a lexicon.
        @param lexicon: the lexicon in Element format
        """
        return self.numEntries(lexicon) / self.numVerbs(lexicon)
    
    def scfs(self, lexicon, verbRegex='.*'):
        """Returns all scfs based on a reg ex of a verb, no doubles
        @param lexicon: the lexicon in Element format
        @param verbRegex: regular expression specifying a class of verbs 
        @rtype: set
        """
        scfs = set([])
        for record in lexicon.getiterator('record'):
            verb = record.find('lx').text
            if re.compile(verbRegex).match(verb):
                scfs.add(record.find('scf').text)
        return scfs    
    
    def numScfs(self, lexicon, verbRegex='.*'):
        """Returns number of scfs based on a reg ex of a verb
        @param lexicon: the lexicon in Element format
        @param verbRegex: regular expression specifying a class of verbs 
        @rtype: integer
        """
        return len(self.scfs(lexicon, verbRegex))
    
    def verbs(self, lexicon, scfRegex='.*'):
        """Returns all verbs based on a reg ex of a scf, no doubles
        @param lexicon: the lexicon in Element format
        @param scfRegex: regular expression specifying a class of scfs 
        @rtype: set
        """
        verbs = set([])
        for record in lexicon.getiterator('record'):
            scf = record.find('scf').text
            if re.compile(scfRegex).match(scf):
                verbs.add(record.find('lx').text)
        return verbs    
    
    def numVerbs(self, lexicon, scfRegex='.*'):
        """Returns number of verbs based on a reg ex of a scf
        @param lexicon: the lexicon in Element format
        @param scfRegex: regular expression specifying a class of scfs 
        @rtype: integer
        """
        return len(self.verbs(lexicon, scfRegex))
    
            
class MyToolboxCorpusErrorChecker (ToolboxCorpusReader):
    """Class that provides the functionality to analyze and existing toolbox file and check it 
    for errors.
    Provides: 
     - Error checking function
    """
    
    def errorCheck(self,lexicon):
        """Call the error finding functions to check syntax error in Toolbox lexicon.
        @param lexicon : lexicon as Element
        @return: lexicon flawed entries
        """
        errorLexicon = XML('<toolbox_data><header /></toolbox_data>') 
        for entry in lexicon.getiterator('record'):
            entry = self.checkFields(entry)
            entry = self.checkScf(entry)
            if not entry.find('err') is None: errorLexicon.append(entry)
        return errorLexicon
    
    def checkFields(self, entry):
        """Check lexicon to see if has exactly one "lx" and one "scf" and 
        at least one "cf" and if not, adds "err" field with appropriate text.
        @param lexicon: entry as Element
        @return: modified entry
        """
        if entry.find("lx") is None: self.__addError(entry,'lx not found;')
        if len(entry.findall("lx")) > 1: self.__addError(entry,'More than one lx found')
        if entry.find("scf") is None: self.__addError(entry,'scf not found;')
        if len(entry.findall("scf")) > 1: self.__addError(entry,'More than one scf found')
        if entry.find("cf") is None: self.__addError(entry, 'cf not found;')
        return entry
    
    def checkScf(self, entry):
        """Checks the scf for errors and adds an err field for found errors.
        @param entry: the entry to be checked as and Element
        """
        for scf in entry.findall('scf'):
            if not self.validScf(scf.text):
                self.__addError(entry, 'Illegal scf')
            else:
                error = self.checkScfFields(scf.text)
                if not error == None: self.__addError(entry, error)
        return entry
        
    def validScf(self, scf):
        """Checks the general syntax of the scf.
        @param scf: the scf to be checked
        """
        return re.compile('^([\w\-\?]+:[\w\-\?]+,)*([\w\-\?]+:[\w\-\?]+)$').match(scf)
              
    def checkScfFields(self, scf):
        """Checks the arguments, only valid combinations are allowed. Adds err field with an error description
        if an error is found.
        @param scf: the scf to be checked
        """        
        Subj = "SUJ:(NP|VPinf|Ssub|VN|\?)"                       
        Obj = "OBJ:(NP|AP|VPinf|VN|Sint|Ssub|\?)"
        DeObj = "DE-OBJ:(VPinf|PP|Ssub|VN|NP|\?)"                   #added NP
        AObj = "A-OBJ:(VPinf|PP|VN|NP|\?)"                          #added NP
        PObj = "P-OBJ:(PP|AdP|VN|NP|VPinf|\?)"                      #added VPinf
        Ato = "ATO:(Srel|PP|AP|NP|VPpart|VPinf|Ssub|XP|\?)"         #added XP
        Ats = "ATS:(NP|PP|AP|AdP|VPinf|Ssub|VPpart|Sint|VN|XP|\?)"  #added XP
        Refl = "refl:CL"                                            #added
        scfRE = re.compile("^("+Subj+"|"+Obj+"|"+DeObj+"|"+AObj+"|"+PObj+"|"+Ato+"|"+Ats+"|"+Refl+")$")
        
        for arg in scf.split(','):
            result = scfRE.match(arg.strip())
            if result == None: return 'Illegal scf argument ' + arg
        
    def __addError(self, entry, msg):
        """Adds and err field to an entry with the appropriate error message.
        @param entry: entry as Element
        @param msg: the error message
        """       
        Error = SubElement(entry, 'err')
        Error.text = msg
            
class MyToolboxCorpusFuser (ToolboxCorpusReader):
    """Class that provides the functionality to merge two toolbox lexicons.
    Provides:
     - Fusing of two lexicons
    """
    __fusedLexicon = {}
       
    def fuse(self, lexicons):
        """Create a new lexicon with the merged data of all input lexicons. An intermediate
        fused lexicon is kept in the dictionary __fusedLexicon. 
        @param lexicon: list of lexicons 
        """
        self.__fusedLexicon = {}
        
        for lexicon in lexicons:
            for record in lexicon.getiterator('record'):
                (myLx, myScf, myData) = self.__parseRecord(record)
                if not myLx in self.__fusedLexicon.keys():
                    self.__addNew(record)
                else:
                    for scf in self.__fusedLexicon[myLx].keys():
                        if self.__merge(record, scf): break 
                    else:
                        self.__addNew(record)                      
        
        return self.__returnToolboxLexicon()
                    
    def __parseRecord(self, record):
        """Gets all data from the record, returns a triple of (lx, scf, [all other fields in a set of tuples]).
        @param record: the record (Element)
        """
        data = set([])
        for field in record.getchildren():
            if field.tag == 'lx': lx = field.text
            elif field.tag == 'scf': scf = field.text
            else: data = data.union(set([(field.tag, field.text)]))
        return (lx, scf, data)
    
    def __addNew(self, record):
        """Adds the record to the lexicon. Overwrites data that is already present.
        @param record: the record (Element)
        """
        (lx, scf, data) = self.__parseRecord(record)
        self.__fusedLexicon[lx] = self.__fusedLexicon.get(lx,{})
        self.__fusedLexicon[lx][scf] = data 
    
    def __merge(self, record, oldScf):
        """Inserts an argument into a existing record. The scf is updated to the new unified
        version. If unification has failed, newScf will be None and the merge operation will
        fail.
        @param record: the record to put into the dictionary
        @param oldScf: the address to put it
        """
        newScf = self.__unifyScfs(oldScf, record.find('scf').text)
        if newScf is None: return False  

        (lx, scf, data) = self.__parseRecord(record)
        self.__fusedLexicon[lx][newScf] = self.__fusedLexicon[lx].get(oldScf,set([])).union(data)
        if not oldScf==newScf: del self.__fusedLexicon[lx][oldScf]
        return True 
    
    def __unifyScfs(self, scf1, scf2):
        """Verifies if two scfs can be unified and returns the new (most specific) scf. If not it returns None.
        @param scf1: first scf
        @param scf2: second scf
        """       
        unifiedScf = []
        scfArgList1 = [a.strip().split(':') for a in scf1.split(',')]
        scfArgList2 = [a.strip().split(':') for a in scf2.split(',')]
        
        for arg1 in scfArgList1:
            if not len(arg1)==2: return None
            for arg2 in scfArgList2:
                if not len(arg2)==2: return None
                match = self.__matchArgs(arg1, arg2)
                if not match is None:
                    scfArgList2.remove(arg2)
                    unifiedScf.append(match)
                    break                   
                                    
        if not (len(scfArgList2)==0 and len(scfArgList1)==len(unifiedScf)): 
            return None
        else: 
            return ','.join(unifiedScf)
    
    def __matchArgs(self, arg1, arg2):       
        """Verifies if two arguments in a scf can be matched.
        @param arg1: pair with synt and POS
        @param arg2: pair with synt and POS
        """
        x = self.__matchString(arg1[0], arg2[0])
        y = self.__matchString(arg1[1], arg2[1])
        if x is None or y is None:
            return None
        else:
            return ':'.join([x,y])
        
    def __matchString(self, arg1, arg2):
        """Tries to match two strings, ? is the wildcard. On failure returns None."""
        if arg1 == arg2: return arg1
        if arg1 == '?': return arg2
        if arg2 == '?': return arg1
        return None
    
    def __returnToolboxLexicon(self):
        """Generates a xml Element with all data that can be handled to put in a file/on screen.
        """
        root = XML('<toolbox_data><header /></toolbox_data>')    
        for lx in self.__fusedLexicon.keys():
            for scf in self.__fusedLexicon[lx].keys(): 
                record = SubElement(root, 'record')
                field = SubElement(record, 'lx')
                field.text = lx
                field = SubElement(record, 'scf')
                field.text = scf
                for (tag, text) in self.__fusedLexicon[lx][scf]:
                    field = SubElement(record, tag)
                    field.text = text        
        return root

class MyExtendedToolboxCorpusReader (MyToolboxCorpusReader, MyToolboxCorpusAnalysis, MyToolboxCorpusErrorChecker, MyToolboxCorpusFuser):
    """The final reader that integrates all tools. Here we can program the functions
    that are directly used by the program. Can access all functions we programmed. 
    """